from typing import List, Union
import cv2
import numpy as np

from ezgl import Matrix4x4

__all__ = [
    'solve_handeye',
]

def solve_handeye(b2g_list: List[Matrix4x4], s2t_list: List[Matrix4x4]):
    """
    Parameters:
    - b2g: base to gripper, gripper's pose in base frame
    - s2t: sensor to target, target's pose in sensor frame

    Returns:
    - g2s: gripper to sensor, sensor's pose in gripper frame
    """
    R_b2g = [A.matrix33 for A in b2g_list]
    t_b2g = [A.xyz for A in b2g_list]
    R_s2t = [B.matrix33 for B in s2t_list]
    t_s2t = [B.xyz  for B in s2t_list]
    
    R_g2s, t_g2s = cv2.calibrateHandEye(
        R_b2g, t_b2g, R_s2t, t_s2t,
        method=cv2.CALIB_HAND_EYE_TSAI
    )
    
    g2s = Matrix4x4.fromRotTrans(R_g2s, t_g2s.squeeze())
    
    # evaluate the result 
    residuals = []
    for i in range(1, len(b2g_list)):
        res = b2g_list[i].inverse() * b2g_list[i-1] * g2s - g2s * s2t_list[i] * s2t_list[i-1].inverse()
        residuals.append(np.linalg.norm(res))
    residual = np.mean(residuals)
    
    return g2s, residual